#include "VideoDecoder.h"


void VideoDecoder::SetImageCallback(DecodedImageCallback func)
{
    _image_callback = func;
}

uint32_t VideoDecoder::ReconnectTimes()
{
    return _reconnect_times;
}

void VideoDecoder::DecodeThread()
{
    AVPacket *av_packet = av_packet_alloc();
    AVFrame *av_frame = av_frame_alloc();
    AVDictionary *options = NULL;
    AVDictionary *codec_opts = NULL;
    uint32_t pkt_idx = 0;

    int64_t latest_pts = 0;
    uint32_t image_idx = 0;

    const AVCodec *codec = NULL;
    AVCodecContext *codec_ctx = NULL;

    AVCodecParameters *codecpar = nullptr;
    int video_index = -1;

    auto hw_device_ctx = SharedResource::GetInstance()->HW_Device_Ctx(_gpu_idx);

    // av_dict_set(&options, "buffer_size", "1024000", 0); //设置缓存大小
    av_dict_set(&options, "buffer_size", "2048000", 0); //设置缓存大小
    // av_dict_set(&options, "rtsp_transport", "tcp", 0); //以tcp的方式打开,
    av_dict_set(&options, "rtsp_transport", "tcp", 0); //以tcp的方式打开,
    av_dict_set(&options, "stimeout", "5000000", 0); //设置超时断开链接时间，单位us
    av_dict_set(&options, "max_delay", "100000", 0); //设置最大时延
    av_dict_set(&options, "fflags", "nobuffer", 0);
    av_dict_set(&options, "flags", "low_delay", 0);
    av_dict_set(&options, "flags2", "fast", 0);

    //申请AVFormatContext
    AVFormatContext *in_fmt_ctx = avformat_alloc_context(); //用来申请AVFormatContext类型变量并初始化默认参数,申请的空间

    //打开网络流或文件流
    if (avformat_open_input(&in_fmt_ctx, _rtsp_url.c_str(), NULL, &options) != 0)
    {
        SPDLOG_INFO("Couldn't open input stream: {} ", _rtsp_url);
        goto finish;
    }

    //获取视频文件信息
    if (avformat_find_stream_info(in_fmt_ctx, NULL) < 0)
    {
        SPDLOG_INFO("Couldn't find stream information: {}", _rtsp_url);
        goto finish;
    }

    //打印输入流的信息
    av_dump_format(in_fmt_ctx, 0, _rtsp_url.c_str(), 0);

    //查找码流中是否有视频流
    for (unsigned i = 0; i < in_fmt_ctx->nb_streams; i++)
    {
        if (in_fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            video_index = i;
            break;
        }
    }
    if (video_index == -1)
    {
        SPDLOG_ERROR("Didn't find a video stream: {}", _rtsp_url);
        goto finish;
    }

    codecpar = in_fmt_ctx->streams[video_index]->codecpar;
    //查找解码器

    if (codecpar->codec_id == AV_CODEC_ID_H264)
    {
#ifdef HWACCELS_INTEL
        codec = avcodec_find_decoder_by_name("h264_qsv");
#elif defined HWACCELS_NVIDIA
        codec = avcodec_find_decoder_by_name("h264_cuvid");
#else
        codec = avcodec_find_decoder(AV_CODEC_ID_H264);
#endif
    }
    else if (codecpar->codec_id == AV_CODEC_ID_HEVC)
    {
#ifdef HWACCELS_INTEL
        codec = avcodec_find_decoder_by_name("hevc_qsv");
#elif defined HWACCELS_NVIDIA
        codec = avcodec_find_decoder_by_name("hevc_cuvid");
#else
        codec = avcodec_find_decoder(AV_CODEC_ID_H264);
#endif
    }

    if (codec == NULL)
    {
        SPDLOG_ERROR("Can't find decoder: {}", (int)codecpar->codec_id);
        goto finish;
    }

    //根据解码方式，申请AVCodecContext
    codec_ctx = avcodec_alloc_context3(codec);

    //将码流参数拷贝到新的AVCodecContext中，并拷贝extradata内容
    if (avcodec_parameters_to_context(codec_ctx, codecpar) < 0)
    {
        SPDLOG_ERROR("avcodec_parameters_to_context failed: {}", _rtsp_url);
        goto finish;
    }

#ifndef HWACCELS_NONE
    if (hw_device_ctx)
    {
        codec_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
    }
    else
    {
        goto finish;
    }

#endif

    // 添加低延迟标志
    codec_ctx->flags |= AV_CODEC_FLAG_LOW_DELAY;
    codec_ctx->thread_count = 4;

    //   这个参数其实没有用，因为相机一般不会有b帧
    codec_ctx->max_b_frames = 0;

    codec_ctx->flags2 |= AV_CODEC_FLAG2_FAST; // 设置 fast 解码
    codec_ctx->delay = 0;

    SPDLOG_WARN("当前使用gpu:{}", _gpu_idx);
    av_dict_set(&codec_opts, "gpu", std::to_string(_gpu_idx).c_str(), 0); // gpu编号

    // 这些参数也没有用
    av_dict_set(&codec_opts, "tune", "zerolatency", 0);
    av_dict_set(&codec_opts, "preset", "ultrafast", 0);

    // 保证gpuMat也使用同样的设备
    cudaSetDevice(_gpu_idx);
    cv::cuda::setDevice(_gpu_idx);

    //初始化AVCodecContext
    if (avcodec_open2(codec_ctx, codec, &codec_opts) < 0)
    {
        SPDLOG_ERROR("avcodec_open2 failed. ");
        goto finish;
    }

    avformat_flush(in_fmt_ctx);

    // 2. 跳转到末尾（比如当前时间戳）
    // latest_pts = av_gettime(); // 当前时间（微秒）
    // av_seek_frame(in_fmt_ctx, -1, latest_pts, AVSEEK_FLAG_BACKWARD | AVSEEK_FLAG_ANY);
    //开始循环读取数据
    while (_start)
    {
        auto ret = av_read_frame(in_fmt_ctx, av_packet);
        if (ret >= 0)
        {
            pkt_idx++;
            if (av_packet->stream_index == video_index)
            {
                avcodec_send_packet(codec_ctx, av_packet);
                if (avcodec_receive_frame(codec_ctx, av_frame) == 0)
                {
                    image_idx++;

                    if (_image_callback)
                    {
                        auto image = std::make_shared<ImageInfo>();
                        image->index = image_idx;
                        image->width = av_frame->width;
                        image->height = av_frame->height;
                        image->time = std::chrono::system_clock::now();
                        cv::cuda::GpuMat rgb(av_frame->height, av_frame->width, CV_8UC3);
                        cv::cuda::GpuMat y_plane(av_frame->height, av_frame->width, CV_8UC1, av_frame->data[0], av_frame->linesize[0]);
                        cv::cuda::GpuMat uv_plane(av_frame->height / 2, av_frame->width, CV_8UC1, av_frame->data[1], av_frame->linesize[1]);

                        const Npp8u *src[2] = { y_plane.ptr<Npp8u>(), uv_plane.ptr<Npp8u>() };
                        int src_step = av_frame->linesize[0];

                        Npp8u *dst = rgb.ptr<Npp8u>();
                        int dst_step = static_cast<int>(rgb.step);

                        NppiSize roi = { av_frame->width, av_frame->height };

                        NppStatus status = nppiNV12ToBGR_8u_P2C3R_Ctx(src, src_step, dst, dst_step, roi, _npp_stream_ctx);

                        if (status != NPP_SUCCESS)
                        {
                            SPDLOG_ERROR("nppiNV12ToBGR_8u_P2C3R_Ctx : {}", (int)status);
                        }

                        // 可以不同主动调用等待同步，在调用ptr,download时会自动同步
                        // cudaStreamSynchronize(ctx.hStream);

                        image->image = rgb;
                        _image_callback(image);
                    }
                }
            }
        }
        else
        {
            // 处理网络断开情况（如EAGAIN、AVERROR_EOF）
            if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
            {
                SPDLOG_INFO("Network disconnection detected, trying to reconnect...");
                goto finish; // 跳转到重连逻辑
            }
        }
        av_packet_unref(av_packet);
    }

finish:
    if (options)
        av_dict_free(&options);

    if (codec_opts)
        av_dict_free(&codec_opts);

    if (codec_ctx)
        avcodec_free_context(&codec_ctx);

    av_frame_free(&av_frame);
    av_packet_free(&av_packet);

    avformat_close_input(&in_fmt_ctx);
}

void VideoDecoder::Init(const std::string &rtsp_url,const int32_t gpu_idx, const int32_t core_idx )
{
    _rtsp_url = rtsp_url;
    _gpu_idx = gpu_idx;
    _core_idx = core_idx;
}


void VideoDecoder::Start()
{
    _start = true;
    if (!_cuda_stream)
    {
        cudaStreamCreate(&_cuda_stream);
    }

    nppGetStreamContext(&_npp_stream_ctx);
    _npp_stream_ctx.hStream = _cuda_stream;

    _thread = std::make_shared<std::thread>([this]() {
        while (_start)
        {
            if (_core_idx > -1)
            {
                cpu_set_t cpuset;
                CPU_ZERO(&cpuset);
                CPU_SET(_core_idx, &cpuset);
                pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
            }

            DecodeThread();

            _reconnect_times++;

            if (_start)
            {
                std::this_thread::sleep_for(3s);
            }
        }
    });
}

void VideoDecoder::Stop()
{
    _start = false;
    if (_thread)
    {
        _thread->join();
        _thread = nullptr;
    }

    if (_cuda_stream)
    {
        cudaStreamDestroy(_cuda_stream);
    }
}

